import numpy as np
import torch
import torch.nn as nn
from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec
from fastervit import create_model


@BACKBONE_REGISTRY.register()
class FasterVit(Backbone):
    def __init__(self, cfg, input_shape):
        super().__init__()
        scale = cfg.MODEL.FASTERVIT.SCALE
        ckpt_path = cfg.MODEL.FASTERVIT.CKPT

        model = create_model(scale, 
                          resolution=[360, 2640],   # 输入size
                          window_size=[7, 7, 12, 6],
                          ct_size=2,
                          dim=64,
                          pretrained=True,
                          model_path=ckpt_path)
        self.patch_embed = model.patch_embed
        self.levels = model.levels

        self._out_features = ["res2", "res3", "res4", "res6"]

        self._out_feature_strides = {
            "res2": 4,
            "res3": 8,
            "res4": 16,
            "res5": 32,
            "res6": 32,
        }
        # self._out_feature_channels = {
        #     "res2": self.num_features[0],
        #     "res3": self.num_features[1],
        #     "res4": self.num_features[2],
        #     "res5": self.num_features[3],
        # }
        self._out_feature_channels = {
            "res2": 64,
            "res3": 128,
            "res4": 256,
            "res5": 512,
            "res6": 512,
        }

    def forward(self, x):
        """
        Args:
            x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
        Returns:
            dict[str->Tensor]: names and the corresponding features
        """
        assert (
            x.dim() == 4
        ), f"FasterVit takes an input of shape (N, C, H, W). Got {x.shape} instead!"
        outputs = {}

        y = {}
        x = self.patch_embed(x)
        y["res2"] = x

        for i, level in enumerate(self.levels):
            res_i = "res{}".format(i+3)
            x = level(x)
            y[res_i] = x

        for k in y.keys():
            if k in self._out_features:
                outputs[k] = y[k]
        return outputs

    def output_shape(self):
        return {
            name: ShapeSpec(
                channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
            )
            for name in self._out_features
        }

    @property
    def size_divisibility(self):
        return 32
